// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/memory/allocation_pool.h>
#include <pollux/common/base/bit_util.h>
#include <pollux/common/base/exceptions.h>
#include <pollux/common/memory/memory_allocator.h>

namespace kumo::pollux::memory {
    melon::Range<char *> AllocationPool::rangeAt(int32_t index) const {
        if (index < allocations_.size()) {
            auto run = allocations_[index].runAt(0);
            return melon::Range<char *>(
                run.data<char>(),
                run.data<char>() == startOfRun_ ? currentOffset_ : run.numBytes());
        }
        const auto largeIndex = index - allocations_.size();
        if (largeIndex < largeAllocations_.size()) {
            auto range = largeAllocations_[largeIndex].hugePageRange().value();
            if (range.data() == startOfRun_) {
                return melon::Range<char *>(range.data(), currentOffset_);
            }
            return range;
        }
        POLLUX_FAIL("Out of range index for rangeAt(): {}", index);
    }

    void AllocationPool::clear() {
        allocations_.clear();
        largeAllocations_.clear();
        startOfRun_ = nullptr;
        bytesInRun_ = 0;
        currentOffset_ = 0;
        usedBytes_ = 0;
    }

    char *AllocationPool::allocateFixed(uint64_t bytes, int32_t alignment) {
        POLLUX_CHECK_GT(bytes, 0, "Cannot allocate zero bytes");
        if (freeAddressableBytes() >= bytes && alignment == 1) {
            auto *result = startOfRun_ + currentOffset_;
            maybeGrowLastAllocation(bytes);
            return result;
        }
        POLLUX_CHECK_EQ(
            __builtin_popcount(alignment), 1, "Alignment can only be power of 2");

        auto numPages = AllocationTraits::numPages(bytes + alignment - 1);

        if (freeAddressableBytes() == 0) {
            newRunImpl(numPages);
        } else {
            auto alignedBytes = bytes + alignmentPadding(firstFreeInRun(), alignment);
            if (freeAddressableBytes() < alignedBytes) {
                newRunImpl(numPages);
            }
        }
        currentOffset_ += alignmentPadding(firstFreeInRun(), alignment);
        POLLUX_CHECK_LE(bytes + currentOffset_, bytesInRun_);
        auto *result = startOfRun_ + currentOffset_;
        POLLUX_CHECK_EQ(reinterpret_cast<uintptr_t>(result) % alignment, 0);
        maybeGrowLastAllocation(bytes);
        return result;
    }

    void AllocationPool::maybeGrowLastAllocation(uint64_t bytesRequested) {
        const auto updateOffset = currentOffset_ + bytesRequested;
        if (updateOffset > endOfReservedRun()) {
            POLLUX_CHECK_GT(bytesInRun_, AllocationTraits::kHugePageSize);
            const auto bytesToReserve = bits::roundUp(
                updateOffset - endOfReservedRun(), AllocationTraits::kHugePageSize);
            largeAllocations_.back().grow(AllocationTraits::numPages(bytesToReserve));
            usedBytes_ += bytesToReserve;
        }
        // Only update currentOffset_ once it points to valid data.
        currentOffset_ = updateOffset;
    }

    void AllocationPool::newRunImpl(MachinePageCount numPages) {
        if (usedBytes_ >= hugePageThreshold_ ||
            numPages > pool_->sizeClasses().back()) {
            // At least 16 huge pages, no more than kMaxMmapBytes. The next is
            // double the previous. Because the previous is a hair under the
            // power of two because of fractional pages at ends of allocation,
            // add an extra huge page size.
            int64_t nextSize = std::min(
                kMaxMmapBytes,
                std::max<int64_t>(
                    16 * AllocationTraits::kHugePageSize,
                    bits::nextPowerOfTwo(
                        usedBytes_ + AllocationTraits::kHugePageSize)));
            // Round 'numPages' to no of pages in huge page. Allocating this plus an
            // extra huge page guarantees that 'numPages' worth of contiguous aligned
            // huge pages will be founfd in the allocation.
            numPages = bits::roundUp(numPages, AllocationTraits::numPagesInHugePage());
            if (AllocationTraits::pageBytes(numPages) +
                AllocationTraits::kHugePageSize >
                nextSize) {
                // Extra large single request.
                nextSize = AllocationTraits::pageBytes(numPages) +
                           AllocationTraits::kHugePageSize;
            }

            ContiguousAllocation largeAlloc;
            const MachinePageCount pagesToAlloc =
                    AllocationTraits::numPagesInHugePage();
            pool_->allocateContiguous(
                pagesToAlloc, largeAlloc, AllocationTraits::numPages(nextSize));

            auto range = largeAlloc.hugePageRange().value();
            startOfRun_ = range.data();
            bytesInRun_ = range.size();
            largeAllocations_.emplace_back(std::move(largeAlloc));
            currentOffset_ = 0;
            usedBytes_ += AllocationTraits::pageBytes(pagesToAlloc);
            return;
        }

        Allocation allocation;
        auto roundedPages = std::max<int32_t>(kMinPages, numPages);
        pool_->allocateNonContiguous(roundedPages, allocation, roundedPages);
        POLLUX_CHECK_EQ(allocation.numRuns(), 1);
        startOfRun_ = allocation.runAt(0).data<char>();
        bytesInRun_ = allocation.runAt(0).numBytes();
        currentOffset_ = 0;
        allocations_.push_back(std::move(allocation));
        usedBytes_ += bytesInRun_;
    }

    void AllocationPool::newRun(int64_t preferredSize) {
        newRunImpl(AllocationTraits::numPages(preferredSize));
    }
} // namespace kumo::pollux::memory
